package main

import (
	"fmt"
	"os"
	"unsafe"

	"github.com/apache/arrow/go/v6/arrow"
	"github.com/apache/arrow/go/v6/arrow/array"
	"github.com/apache/arrow/go/v6/arrow/ipc"
	"github.com/apache/arrow/go/v6/arrow/memory"
)

func write(file_path string) {
	file, err := os.Create(file_path)
	// file, err := os.OpenFile(file_path, os.O_WRONLY, 0600)
	if err != nil {
		fmt.Errorf("create file {%s} error: %v", file_path, err)
	}
	defer file.Close()

	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
	record := init_record_slice(mem)

	// 目前仅支持如下两种压缩算法
	// zstd 压缩
	// ipc.WithZstd()
	// lz4压缩
	// ipc.WithLZ4()
	writer, err := ipc.NewFileWriter(file, ipc.WithSchema(record.Schema()), ipc.WithAllocator(mem), ipc.WithZstd())
	if err != nil {
		fmt.Errorf("create writer of file {%s} error: %v", file_path, err)
	}
	defer writer.Close()

	schema := record.Schema()
	fmt.Printf("record slice rows %d cols %d\n", record.NumRows(), record.NumCols())
	for i, col := range record.Columns() {
		fmt.Printf("field: %v type: %v rows: %v\t", schema.Field(i).Name, col.Data().DataType(), col.Len())
		fmt.Printf("type: %v rows: %v\t", col.Data().DataType(), col.Len())
		fmt.Printf("col: %v\n", col)
	}
	errw := writer.Write(record)
	if errw != nil {
		fmt.Printf("write error: %v\n", errw)
	}
}

func read(file_path string) {
	file, err := os.OpenFile(file_path, os.O_RDONLY, 0400)
	if err != nil {
		fmt.Errorf("open file {%s} error: %v", file_path, err)
	}
	defer file.Close()

	reader, err := ipc.NewFileReader(file, ipc.WithZstd())
	if err != nil {
		fmt.Errorf("create reader of file {%s} error: %v", file_path, err)
	}
	defer reader.Close()

	fmt.Printf("records: %d\n", reader.NumRecords())

	record, err := reader.Record(0)
	if err != nil {
		fmt.Errorf("get record from reader error %v\n", err)
	}

	schema := record.Schema()
	fmt.Printf("record slice rows %d cols %d\n", record.NumRows(), record.NumCols())
	for i, col := range record.Columns() {
		fmt.Printf("field: %v type: %v rows: %v\t", schema.Field(i).Name, col.Data().DataType(), col.Len())
		fmt.Printf("col: %v\n", col)
	}
	// reader.Schema()
	_ = reader
}

func init_table_from_array() *array.Table {
	// 定义数据行数目
	const rows = 5
	// 定义数据列
	id_field := arrow.Field{Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}
	name_field := arrow.Field{Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}
	age_field := arrow.Field{Name: "age", Type: arrow.PrimitiveTypes.Int8, Nullable: true}
	sex_field := arrow.Field{Name: "sex", Type: &arrow.BooleanType{}, Nullable: true}
	// 定义当前数据表的schema(id int64, name string, age int8, sex  bool)
	schema := arrow.NewSchema(
		[]arrow.Field{
			id_field, name_field, age_field, sex_field,
		}, nil,
	)

	alloc := memory.NewCheckedAllocator(memory.NewGoAllocator())

	valid := []bool{true, true, true, true, true}
	id_builder := array.NewInt64Builder(alloc)
	id_builder.AppendValues([]int64{1, 2, 3, 4, 5}, valid)
	id_val := id_builder.NewInt64Array()
	id_chunked := array.NewChunked(arrow.PrimitiveTypes.Int64, []array.Interface{id_val})

	name_builder := array.NewStringBuilder(alloc)
	name_builder.AppendValues([]string{"lynn yuan", "yuan xl", "yuan yy", "yuan xf", "zhou yk"}, valid)
	name_val := name_builder.NewStringArray()
	name_chunked := array.NewChunked(arrow.BinaryTypes.String, []array.Interface{name_val})

	age_builder := array.NewInt8Builder(alloc)
	age_builder.AppendValues([]int8{32, 32, 31, 31, 32}, valid)
	age_val := age_builder.NewInt8Array()
	age_chunked := array.NewChunked(arrow.PrimitiveTypes.Int8, []array.Interface{age_val})

	sex_builder := array.NewBooleanBuilder(alloc)
	sex_builder.AppendValues([]bool{true, true, false, false, true}, valid)
	sex_val := sex_builder.NewBooleanArray()
	sex_chunked := array.NewChunked(&arrow.BooleanType{}, []array.Interface{sex_val})

	// 定义数据列
	columns := []array.Column{
		*array.NewColumn(id_field, id_chunked),
		*array.NewColumn(name_field, name_chunked),
		*array.NewColumn(age_field, age_chunked),
		*array.NewColumn(sex_field, sex_chunked),
	}

	// 定义数据表
	table := array.NewTable(schema, columns, rows)

	fmt.Printf("table rows %d cols %d\n", table.NumRows(), table.NumCols())

	return (*array.Table)(unsafe.Pointer(table))
}

func init_record_slice(mem memory.Allocator) array.Record {
	// 定义数据行数目
	const rows = 5
	// 定义数据列
	id_field := arrow.Field{Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: false}
	name_field := arrow.Field{Name: "name", Type: arrow.BinaryTypes.String, Nullable: true}
	age_field := arrow.Field{Name: "age", Type: arrow.PrimitiveTypes.Int8, Nullable: true}
	sex_field := arrow.Field{Name: "sex", Type: &arrow.BooleanType{}, Nullable: true}
	// 定义当前数据表的schema(id int64, name string, age int8, sex  bool)
	schema := arrow.NewSchema(
		[]arrow.Field{
			id_field, name_field, age_field, sex_field,
		}, nil,
	)

	record_builder := array.NewRecordBuilder(mem, schema)
	// defer record_builder.Release()
	record_builder.Retain()

	record_builder.Field(0).(*array.Int64Builder).AppendValues([]int64{1, 2, 3, 4, 5}, nil)
	record_builder.Field(1).(*array.StringBuilder).AppendValues([]string{"lynn yuan", "yuan xl", "yuan yy", "yuan xf", "zhou yk"}, nil)
	record_builder.Field(2).(*array.Int8Builder).AppendValues([]int8{32, 32, 31, 31, 32}, nil)
	record_builder.Field(3).(*array.BooleanBuilder).AppendValues([]bool{true, true, false, false, true}, nil)

	rec := record_builder.NewRecord()
	rec.Retain()
	// defer rec.Release()

	// rec_slice := rec.NewSlice(0, 4)
	// _ = rec_slice

	fmt.Printf("record slice rows %d cols %d\n", rec.NumRows(), rec.NumCols())

	return rec
}

func main() {
	// table := init_table_from_array()
	// _ = table

	file_path := "/tmp/user.aw"
	write(file_path)

	fmt.Printf("-------------------Reader %s-------------------\n", file_path)

	read(file_path)
}
